#include <iostream>
#include <algorithm>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

typedef struct Ellipse {
    float centerX, centerY;
    float major, minor, angle;
    float score;
} Ellipse;

static bool _find_center(const std::vector<cv::Point>& edgeCoords, const cv::Mat& edges, int r1, int r2, int r3, int neighbor, float& centerX, float& centerY) {
    std::vector<cv::Point> pts;
    pts.push_back(edgeCoords[r1]);
    pts.push_back(edgeCoords[r2]);
    pts.push_back(edgeCoords[r3]);
    
    std::vector<float> ks, bs;
    for (int i = 0; i < pts.size(); i++) {
        int xstart = pts[i].x - neighbor / 2;
        int ystart = pts[i].y - neighbor/ 2;
        std::vector<cv::Point> fit_pts;
        for (int row = ystart; row < ystart + neighbor; row++) {
            for (int col = xstart; col < xstart + neighbor; col++) {
                if (row >= 0 && row < edges.rows && col >= 0 && col < edges.cols) {
                    if (edges.ptr<uchar>(row)[col] > 0) {
                        fit_pts.push_back(cv::Point(col, row));
                    }
                }
            }
        }
        
        cv::Vec4f line;
        cv::fitLine(fit_pts, line, cv::DIST_L2, 0, 0.001, 0.001);
        if (std::abs(line[0]) < FLT_EPSILON) {
            return false;
        } else {
            float k = line[1] / line[0];
            ks.push_back(k);
            bs.push_back(pts[i].y - k*pts[i].x);
        }
    }
    
    std::vector<std::pair<int, int>> lines;
    lines.push_back(std::make_pair<int, int>(0, 1));
    lines.push_back(std::make_pair<int, int>(1, 2));
    std::vector<cv::Point> intersects;
    std::vector<float> slopes, intercepts;
    for (int i = 0; i < lines.size(); i++) {
        float k1 = ks[lines[i].first], k2 = ks[lines[i].second];
        float b1 = bs[lines[i].first], b2 = bs[lines[i].second];
        if (std::abs(k1-k2) < FLT_EPSILON) {
            return false;
        }
        float intersectX = (b2-b1)/(k1-k2);
        float intersectY = k1*intersectX + b1;
        intersects.push_back(cv::Point(intersectX, intersectY));
        
        float midX = (pts[lines[i].first].x + pts[lines[i].second].x) / 2.0;
        float midY = (pts[lines[i].first].y + pts[lines[i].second].y) / 2.0;
        if (std::abs(midX - intersectX) < FLT_EPSILON) {
            return false;
        }
        float slope = (midY - intersectY) / (midX - intersectX);
        float intercept = intersectY - slope * intersectX;
        slopes.push_back(slope);
        intercepts.push_back(intercept);
    }
    
    float k1 = slopes[0], k2 = slopes[1];
    float b1 = intercepts[0], b2 = intercepts[1];
    if (std::abs(k1-k2) < FLT_EPSILON) {
        return false;
    }
    centerX = (b2 - b1) / (k1-k2);
    centerY = k1 * centerX + b1;
    cv::Rect rect = cv::boundingRect(edgeCoords);
    float left = rect.x + rect.width / 3.0;
    float right = rect.x + rect.width * 2 / 3.0;
    float top = rect.y + rect.height / 3.0;
    float bottom  =rect.y + rect.height* 2/ 3.0;
    if (centerX > left && centerX < right && centerY > top && centerY < bottom) {
        return true;
    }
    return false;
}

static bool _find_semi_axis(const std::vector<cv::Point>& edgeCoords, int r1, int r2, int r3, float centerX, float centerY, float& major, float& minor, float& angle) {
    std::vector<cv::Point> pts;
    pts.push_back(edgeCoords[r1]);
    pts.push_back(edgeCoords[r2]);
    pts.push_back(edgeCoords[r3]);
    
    float x1 = edgeCoords[r1].x - centerX;
    float y1 = edgeCoords[r1].y - centerY;
    float x2 = edgeCoords[r2].x - centerX;
    float y2 = edgeCoords[r2].y - centerY;
    float x3 = edgeCoords[r3].x - centerX;
    float y3 = edgeCoords[r3].y - centerY;
    
    float CpieDenominator = (x1*x1*x2*y2 - x2*x2*x1*y1) * (x3*x3*y2*y2 - x2*x2*y3*y3) - (x2*x2*x3*y3 - x3*x3*x2*y2) * (x2*x2*y1*y1 - x1*x1*y2*y2);
    float BpieDenominator = x1*x1*x2*y2 - x2*x2*x1*y1;
    float ApieDenominator = x1*x1;
    if (std::abs(ApieDenominator) < FLT_EPSILON || std::abs(BpieDenominator) < FLT_EPSILON || std::abs(CpieDenominator) < FLT_EPSILON) {
        return false;
    }
    float CpieNumerator = (x1*x1 - x2*x2) * (x2*x2*x3*y3 - x3*x3*x2*y2) - (x2*x2 - x3*x3) * (x1*x1*x2*y2 - x2*x2*x1*y1);
    float Cpie = CpieNumerator / CpieDenominator;
    float BpieNumerator = (x2*x2*y1*y1 - x1*x1*y2*y2) * Cpie - (x1*x1 - x2*x2);
    float Bpie = BpieNumerator / BpieDenominator;
    float ApieNumerator = 1 - x1*y1*Bpie - y1*y1*Cpie;
    float Apie = ApieNumerator / ApieDenominator;
    if (Bpie * Bpie - 4*Apie*Cpie >= 0) {
        return false;
    }
    
    float theta = 0;
    if (std::abs(Bpie) < FLT_EPSILON) {
        if (Apie < Cpie) {
            theta = 0;
        } else {
            theta = CV_PI / 2;
        }
    } else {
        theta = std::atan((Cpie - Apie - std::sqrt(std::pow(Apie-Cpie, 2) + Bpie*Bpie)) / Bpie);
    }
    
    float denom1 = Apie*std::pow(std::sin(theta), 2) - Cpie * std::pow(std::cos(theta), 2);
    if (std::abs(denom1) < FLT_EPSILON) {
        return false;
    }
    float b = std::sqrt((std::pow(std::sin(theta), 4) - std::pow(std::cos(theta), 4)) / denom1);
    float denom2 = Apie - std::pow(std::sin(theta), 2) / (b*b);
    if (std::abs(denom2) < FLT_EPSILON) {
        return false;
    }
    float a = std::sqrt(std::pow(std::cos(theta), 2) / denom2);
    if (a > 0 && b > 0) {
        angle = theta;
        major = a;
        minor = b;
        return true;
    }
    return false;
}

static int _is_similar(const std::vector<Ellipse>& ellipses, const std::vector<cv::Point>& contour, float centerX, float centerY, float major, float minor, float angle) {
    int idx = -1;
    cv::Rect rect = cv::boundingRect(contour);
    float deltaMin = std::min(rect.width, rect.height);
    float deltaMax = std::max(rect.width, rect.height);
    for (int i = 0; i < ellipses.size(); i++) {
        float centerXDiff = std::abs(ellipses[i].centerX - centerX);
        float centerYDiff = std::abs(ellipses[i].centerY - centerY);
        float centerDiff = std::sqrt(centerXDiff*centerXDiff + centerYDiff*centerYDiff);
        float majorDiff = std::abs(ellipses[i].major - major);
        float minorDiff = std::abs(ellipses[i].minor - minor);
        if (centerDiff < deltaMin && major < deltaMax && minor < deltaMin) {
            return i;
        }
    }
    return idx;
}

static float _average_weight(float old, float now, float score) {
    return (old*score+now) / (score+1);
}

static int _compare_score(const Ellipse& e1, const Ellipse& e2) {
    if (e1.score > e2.score) {
        return 1;
    }
    if (e1.score < e2.score) {
        return -1;
    }
    return 0;
}

static bool _randomized_hough_ellipse(cv::Mat& dst, const cv::Mat& edges, const std::vector<cv::Point>& contour, int maxIter, Ellipse& ellipse) {
    CV_Assert(edges.type() == CV_8UC1);
    int contourSize = contour.size();
    if (contourSize < 15) {
        return false;
    }
    
    std::vector<Ellipse> ellipses;
    cv::RNG rng = cv::RNG(10);
    for (int i = 0; i < maxIter; i++) {
        int r1 = rng.uniform(0, contourSize);
        int r2 = rng.uniform(0, contourSize);
        int r3 = rng.uniform(0, contourSize);
        int neighbor = 11;
        float centerX, centerY;
        if (!_find_center(contour, edges, r1, r2, r3, neighbor, centerX, centerY)) {
            continue;
        }
        //cv::circle(dst, cv::Point(cvRound(centerX), cvRound(centerY)), 3, cv::Scalar(0, 0, 255), cv::FILLED, cv::LINE_AA);
        //std::cout << "center = " << centerX << ", " << centerY << std::endl;
        
        float major, minor, angle;
        if (!_find_semi_axis(contour, r1, r2, r3, centerX, centerY, major, minor, angle)) {
            continue;
        }
        
        int idx = _is_similar(ellipses, contour, centerX, centerY, major, minor, angle);
        if (idx == -1) {
            Ellipse ellipse;
            ellipse.centerX = centerX;
            ellipse.centerY = centerY;
            ellipse.major = major;
            ellipse.minor = minor;
            ellipse.angle = angle;
            ellipse.score = 1;
            ellipses.push_back(ellipse);
        } else {
            ellipses[idx].score += 1;
            float score = ellipses[idx].score;
            ellipses[idx].centerX = _average_weight(ellipses[idx].centerX, centerX, score);
            ellipses[idx].centerY = _average_weight(ellipses[idx].centerY, centerY, score);
            ellipses[idx].major = _average_weight(ellipses[idx].major, major, score);
            ellipses[idx].minor = _average_weight(ellipses[idx].minor, minor, score);
            ellipses[idx].angle = _average_weight(ellipses[idx].angle, angle, score);
        }
    }
    
    if (ellipses.size() > 0) {
        std::sort(ellipses.begin(), ellipses.end(), _compare_score);
        ellipse.centerX = ellipses[0].centerX;
        ellipse.centerY = ellipses[0].centerY;
        ellipse.major = ellipses[0].major;
        ellipse.minor = ellipses[0].minor;
        ellipse.angle = ellipses[0].angle;
        ellipse.score = ellipses[0].score;
        return true;
    }
    
    return false;
}

int main(int argc, char **argv) {
    cv::Mat src = cv::imread("/home/xlll/Downloads/opencv/samples/data/detect_blob.png", cv::IMREAD_COLOR);
    if (src.empty()) {
        std::cout << "Failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat gray;
    cv::cvtColor(src, gray, cv::COLOR_BGR2GRAY);
    int kernelSize = 3;
    double cannyThresh = 50;
    cv::Mat edges, dx, dy;
    cv::Sobel(gray, dx, CV_16SC1, 1, 0, kernelSize, 1, 0, cv::BORDER_REPLICATE);
    cv::Sobel(gray, dy, CV_16SC1, 0, 1, kernelSize, 1, 0, cv::BORDER_REPLICATE);
    cv::Canny(dx, dy, edges, cannyThresh/2, cannyThresh, false);
    
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(edges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);
    
    cv::Mat dst = src.clone();
    int maxIter = 10;
    //for (int i = 7; i < 8; i++) {
    for (int i = 0; i < contours.size(); i++) {
        Ellipse ellipse;
        if (_randomized_hough_ellipse(dst, edges, contours[i], maxIter, ellipse)) {
            cv::ellipse(dst, cv::Point(cvRound(ellipse.centerX), cvRound(ellipse.centerY)), cv::Size(cvRound(ellipse.major), cvRound(ellipse.minor)), ellipse.angle*180/CV_PI, 0, 360, cv::Scalar(255, 0, 255), 2, cv::LINE_AA);
            std::cout << ellipse.score << std::endl;
        }
    }
    
    cv::imshow("src", src);
    cv::imshow("edges", edges);
    cv::imshow("dst", dst);
    cv::waitKey(0);
    
    return 0;
}
